from PyCmpltrtok.common import sep, get_dir_name_ext, uuid
import os
import json
from python_nlp.datasets.load_large_json import load_large_json

NAME2QWEN_ROLE = {
    '病人：': 'user',
    '医生：': 'assistant',
}


def med_dialog_list2qwen_format(xlist):
    obj = {
        'id': uuid(),
        'conversations': []
    }
    for element in xlist:
        for key in NAME2QWEN_ROLE:
            xlen_key = len(key)
            if element[:xlen_key] == key:
                obj['conversations'].append({
                    'from': NAME2QWEN_ROLE[key],
                    'value': element[xlen_key:]
                })
    return obj


def med_dialog_make_qwen_json(xfile_path, n_turn=0, num=6):

    print(f'Loading {xfile_path}')
    xdir, xbase, xext = get_dir_name_ext(xfile_path)
    if n_turn:
        n_turn_str = f'-{n_turn}turn'
    else:
        n_turn_str = ''
    if num:
        num_str = f'-x{num}'
    else:
        num_str = ''
    dataset_path = os.path.join(xdir, f'{xbase}.qwen{n_turn_str}{num_str}.json')
    print(f'Writing to {dataset_path}')

    with open(dataset_path, 'w', encoding='utf8') as f:
        f.write('[\n')
        first = True
        i = 0
        cnt = 0
        for x in load_large_json(xfile_path):
            i += 1
            xlen = len(x)
            if not n_turn or xlen == n_turn * 2:
                if num and num <= 10:
                    sep()
                    print(i, x)
                    sep()
                    print('\n'.join(x))
                    sep()
                elif cnt % 1000 == 0:
                    print(cnt, end=', ', flush=True)
                if not first:
                    f.write(',')
                else:
                    first = False
                source = med_dialog_list2qwen_format(x)
                source_str = json.dumps(source, ensure_ascii=False)
                f.write(source_str + '\n')
                cnt += 1
                if num and cnt >= num:
                    break
        f.write(']\n')
        if not num or num > 10:
            print(cnt, flush=True)
            
    print('Done')
    
if __name__ == '__main__':

    def main():
        xpaths = []
        # xpaths.append('D:\_dell7590_root\local\LNP_datasets\med\med_dialog\MedDialog_processed\MedDialog_processed\validate_data.json')
        xpaths.append('/mnt/d/_const/large_data/NLP/med/med_dialog/MedDialog_processed/MedDialog_processed/validate_data.json')
        # xpaths.append('/home/yunpeng/datasets/med_dialog/train_data.json')
        # xpaths.append('/home/yunpeng/datasets/med_dialog/validate_data.json')
        # xpaths.append('/home/yunpeng/datasets/med_dialog/test_data.json')
        for p in xpaths:
            # med_dialog_make_qwen_json(p, num=4, n_turn=3)
            med_dialog_make_qwen_json(p, num=20, n_turn=0)
            # med_dialog_make_qwen_json(p, num=0, n_turn=0)

    main()
